# -*- coding: utf-8 -*-

import os
import numpy as np
import data5.utils as utils
import data5.config as config


def default_get_char_label(label, index):
    return label[index]


def generator_char_image_and_char_label(path, image_suffix, mapping_line_parser, image_preprocessor, cut_image_func,
                                        mapping_filename='mappings.txt',
                                        loop_times=1, get_char_label=default_get_char_label):
    """
    创建一个生成器，该生成器不断生产 mapping_filename 中的验证码的灰度图像中的每一个字符图像和该验证码字符的 label
    :param path: mappings.txt 文件和验证码文件所在的路径
    :param image_suffix: 图片的后缀名
    :param mapping_line_parser: 解析 mappings.txt 每一行字符串的函数，该函数参数为 mapping.txt 的一行字符串，该函数返回验证码图片的文件名和该验证码的 label
    :param image_preprocessor: 处理验证码灰度图像的函数，将验证码灰度图像拆成每个字符返回，该函数参数为验证码的灰度图像，该函数返回一个列表保存每一个字符的图像
    :param mapping_filename: mapping 文件的文件名，默认为 mappings.txt
    :param loop_times: mapping_filename 循环读取的次数，如果该参数小于 0 表示无限循环
    :param get_char_label: 获取每个char的label
    :return: None
    """
    full_gen = generator_full_image_and_label(path=path, image_suffix=image_suffix,
                                              mapping_line_parser=mapping_line_parser,
                                              mapping_filename=mapping_filename,
                                              loop_times=loop_times)
    for full_yzm_img, yzm_items, label_yzm in full_gen:
        preprocessed_yzm_items = []
        for yzm_item in yzm_items:
            preprocessed_yzm_items.append(image_preprocessor(yzm_item))

        for index, char_image in enumerate(cut_image_func(image_preprocessor(full_yzm_img))):
            for i in range(9):
                yield char_image, preprocessed_yzm_items[i], str(i) == get_char_label(label_yzm, index)


def generator_full_image_and_label(path, image_suffix, mapping_line_parser, mapping_filename='mappings.txt',
                                   loop_times=1):
    """
    创建一个生成器，该生成器不断生产 mapping_filename 中的验证码的灰度图像中的每一个字符图像和该验证码字符的 label
    :param path: mappings.txt 文件和验证码文件所在的路径
    :param image_suffix: 图片的后缀名
    :param mapping_line_parser: 解析 mappings.txt 每一行字符串的函数，该函数参数为 mapping.txt 的一行字符串，该函数返回验证码图片的文件名和该验证码的 label
    :param image_preprocessor: 处理验证码灰度图像的函数，将验证码灰度图像拆成每个字符返回，该函数参数为验证码的灰度图像，该函数返回一个列表保存每一个字符的图像
    :param mapping_filename: mapping 文件的文件名，默认为 mappings.txt
    :param loop_times: mapping_filename 循环读取的次数，如果该参数小于 0 表示无限循环
    :return: None
    """
    with open(os.path.join(path, mapping_filename), 'r') as f:
        while True:
            if loop_times == 0:
                break
            if loop_times > 0:
                loop_times -= 1

            f.seek(0)

            for line in f.readlines():
                filename, label_yzm = mapping_line_parser(line.strip())

                full_yzm_img = utils.load_image_by_filename(
                    os.path.join(path, filename + '/' + filename + '.' + image_suffix))
                yzm_items = []
                for i in range(9):
                    yzm_items.append(
                        utils.load_image_by_filename(os.path.join(path, filename + '/' + str(i) + '.' + image_suffix)))
                yield full_yzm_img, yzm_items, label_yzm


def generator_next_batch(path, image_suffix, mapping_line_parser, image_preprocessor, cut_image_func,
                         batch_size=50,
                         mapping_filename='mappings.txt', get_char_label=default_get_char_label):
    """
    创建一个生成器，该生成器不断批次生产 mapping_filename 中的验证码的灰度图像中的每一个字符图像和该验证码字符的 label
    :param path: mappings.txt 文件和验证码文件所在的路径
    :param image_suffix: 图片的后缀名
    :param mapping_line_parser: 解析 mappings.txt 每一行字符串的函数，该函数参数为 mapping.txt 的一行字符串，该函数返回验证码图片的文件名和该验证码的 label
    :param image_preprocessor: 处理验证码灰度图像的函数，将验证码灰度图像拆成每个字符返回，该函数参数为验证码的灰度图像，该函数返回一个列表保存每一个字符的图像
    :param char2one_hot_func: 将字符转为 one hot 的函数
    :param batch_size 每批数据的大小
    :param mapping_filename: mapping 文件的文件名，默认为 mappings.txt
    :param get_char_label: 获取每个char的label
    :return: None
    """
    gen = generator_char_image_and_char_label(path=path, image_suffix=image_suffix,
                                              mapping_line_parser=mapping_line_parser,
                                              image_preprocessor=image_preprocessor, cut_image_func=cut_image_func,
                                              mapping_filename=mapping_filename,
                                              loop_times=-1, get_char_label=get_char_label)
    current = 0
    char_images = []
    char_yzm_items = []
    char_labels = []
    for char_image, yzm_items, char_label in gen:
        char_images.append(utils.image_padding(char_image, config.char_image_width, config.char_image_height))
        char_yzm_items.append(yzm_items)
        char_labels.append(char_label)
        current = current + 1
        if current >= batch_size:
            yield np.array(char_images), np.array(char_yzm_items), np.array(char_labels)
            current = 0
            char_images = []
            char_yzm_items = []
            char_labels = []
